package com.aiwiown.snackmq.broker;

import com.aiwiown.snackmq.broker.config.BrokerConfig;
import com.aiwiown.snackmq.cluster.raft.SnackMQStateMachine;
import com.aiwiown.snackmq.cluster.raft.manager.SnackMQRaftNodeManager;
import org.awaitility.core.ConditionTimeoutException;

import java.io.File;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;

import static org.awaitility.Awaitility.await;

/**
 * 一个用于手动启动 SnackMQ Broker 集群的工具类。
 * 运行此 main 方法将启动一个包含 3 个节点的集群，并使其保持运行状态，直到手动终止。
 * 这对于本地开发和与 Console 联调非常有用。
 */
public class ClusterRunner {

    private static final int CLUSTER_SIZE = 3;
    private static final String HOST = "localhost";
    private static final String BASE_DATA_PATH = "data/manual-cluster";
    // 【修复】定义固定的基础端口，使集群地址可预测
    private static final int BASE_SERVER_PORT = 9876;
    private static final int BASE_RAFT_PORT = 9877;

    public static void main(String[] args) throws Exception {
        System.out.println("--- 准备启动 SnackMQ " + CLUSTER_SIZE + "-节点集群 ---");

        // 1. 清理旧的数据目录（可选，确保每次启动都是全新的）
        File baseDir = new File(BASE_DATA_PATH);
        if (baseDir.exists()) {
            System.out.println("清理旧的数据目录: " + baseDir.getAbsolutePath());
            deleteDirectory(baseDir);
        }

        final List<BrokerServer> cluster = new ArrayList<>();
        final List<BrokerConfig> configs = new ArrayList<>();
        final List<String> raftAddresses = new ArrayList<>();
        final Map<String, Integer> raftToBusinessPort = new HashMap<>();

        // 2. 为每个 Broker 生成独立的、固定的配置
        System.out.println("--- 生成固定的 Broker 配置 ---");
        for (int i = 0; i < CLUSTER_SIZE; i++) {
            BrokerConfig config = new BrokerConfig();
            config.setHost(HOST);
            // 【修复】使用固定的、递增的端口号，而不是随机端口
            config.setServerPort(BASE_SERVER_PORT + (i * 2)); // e.g., 9876, 9878, 9880
            config.setRaftPort(BASE_RAFT_PORT + (i * 2));   // e.g., 9877, 9879, 9881

            config.setStorePath(BASE_DATA_PATH + "/broker-" + i + "/store");
            config.setRaftDataPath(BASE_DATA_PATH + "/broker-" + i + "/raft");
            config.setRaftElectionTimeoutMs(3000);
            config.setRaftGroupId("snackmq-manual-cluster");

            // 启用认证，以便 Console 可以连接
            config.setAuthEnabled(true);
            // 使用默认的 admin/admin 凭据
            config.setDefaultAdminUser("admin");
            config.setDefaultAdminPassword("admin");

            configs.add(config);
            String raftAddr = config.getRaftAddress();
            raftAddresses.add(raftAddr);
            raftToBusinessPort.put(raftAddr, config.getServerPort());
            System.out.printf("  - Broker %d: Server Port=%d, Raft Port=%d%n", i, config.getServerPort(), config.getRaftPort());
        }

        // 3. 设置集群的 Raft 初始配置
        String raftInitialConf = String.join(",", raftAddresses);
        for (BrokerConfig config : configs) {
            config.setRaftInitialConf(raftInitialConf);
            config.setRaftAddressToBusinessPortMap(raftToBusinessPort);
        }

        // 4. 创建并启动所有 Broker 实例
        cluster.addAll(configs.stream().map(BrokerServer::new).collect(Collectors.toList()));
        System.out.println("--- 开始启动 Broker 节点... ---");
        cluster.parallelStream().forEach(broker -> {
            try {
                broker.start();
                System.out.println("Broker 已在端口 " + broker.getBrokerConfig().getServerPort() + " (Raft: " + broker.getBrokerConfig().getRaftPort() + ") 启动。");
            } catch (Exception e) {
                System.err.println("Broker 启动失败: " + e.getMessage());
                e.printStackTrace();
            }
        });

        // 5. 【最终修复】等待集群选举出 Leader 并完成状态恢复
        System.out.println("--- 等待集群选举 Leader 并完成状态恢复... ---");
        try {
            // 5.1 等待元数据 Leader
            BrokerServer metaLeader = findMetaLeader(cluster, 30);
            System.out.println("--- 元数据 Leader 已选出: " + metaLeader.getBrokerConfig().getBrokerAddress() + " ---");

            // 5.2 等待数据 Leader
            BrokerServer dataLeader = findDataLeader(cluster, 30);
            System.out.println("--- 数据 Leader 已选出: " + dataLeader.getBrokerConfig().getBrokerAddress() + " ---");

            // 5.3 【关键】等待数据 Leader 的状态机完成恢复
            await("Data state machine on leader " + dataLeader.getBrokerConfig().getServerPort() + " failed to recover")
                    .atMost(30, TimeUnit.SECONDS) // 给予恢复足够的时间
                    .pollInterval(500, TimeUnit.MILLISECONDS)
                    .until(() -> {
                        SnackMQRaftNodeManager nodeManager = dataLeader.getDataRaftNodeManager();
                        if (nodeManager == null) return false;
                        SnackMQStateMachine stateMachine = nodeManager.getStateMachine();
                        // 检查状态机是否已不再处于恢复模式
                        return stateMachine != null && !stateMachine.isRecovering();
                    });
            System.out.println("--- 数据状态机已完成恢复 ---");


            System.out.println("\n=======================================================");
            System.out.println(" SnackMQ 集群已成功启动并运行！");
            System.out.println(" 你现在可以将 Console 或客户端连接到以下任一固定地址:");
            configs.forEach(c -> System.out.println(" -> " + c.getBrokerAddress()));
            System.out.println(" Console 默认管理员凭据: admin / admin");
            System.out.println(" 按 Ctrl+C 停止集群。");
            System.out.println("=======================================================\n");

        } catch (ConditionTimeoutException e) {
            System.err.println("!!! 集群在规定时间内未能选举出 Leader 或完成恢复，请检查日志。 !!!");
            e.printStackTrace();
            // 终止所有已启动的 broker
            cluster.forEach(BrokerServer::shutdown);
            return;
        }

        // 6. 阻塞主线程，使 Broker 持续运行
        new CountDownLatch(1).await();
    }

    private static BrokerServer findMetaLeader(List<BrokerServer> brokers, int timeoutSeconds) {
        return findLeader(brokers, b -> b.getMetadataManager() != null && b.getMetadataManager().isLeader(), "元数据", timeoutSeconds);
    }

    private static BrokerServer findDataLeader(List<BrokerServer> brokers, int timeoutSeconds) {
        return findLeader(brokers, b -> b.getDataRaftNodeManager() != null && b.getDataRaftNodeManager().isLeader(), "数据", timeoutSeconds);
    }

    private static BrokerServer findLeader(List<BrokerServer> brokers, java.util.function.Predicate<BrokerServer> leaderCheck, String leaderType, int timeoutSeconds) {
        AtomicReference<BrokerServer> leaderRef = new AtomicReference<>();
        try {
            await().atMost(timeoutSeconds, TimeUnit.SECONDS)
                    .pollInterval(500, TimeUnit.MILLISECONDS)
                    .until(() -> brokers.stream()
                            .filter(b -> b != null && b.isStarted())
                            .filter(leaderCheck)
                            .findFirst()
                            .map(leader -> {
                                leaderRef.set(leader);
                                return true;
                            }).orElse(false));
        } catch (ConditionTimeoutException e) {
            System.err.printf("在 %d 秒内未能找到 %s Leader%n", timeoutSeconds, leaderType);
            throw e;
        }
        return leaderRef.get();
    }

    private static void deleteDirectory(File directory) {
        File[] allContents = directory.listFiles();
        if (allContents != null) {
            for (File file : allContents) {
                deleteDirectory(file);
            }
        }
        directory.delete();
    }
}